package com.hhq.cloud.websocket;

import io.fabric8.kubernetes.client.KubernetesClient;
import io.fabric8.kubernetes.client.dsl.ExecWatch;
import org.apache.tomcat.util.threads.ThreadPoolExecutor;
import org.springframework.util.MultiValueMap;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class K8sTerminalWebSocketHandler extends TextWebSocketHandler {

    private final Map<String, ExecWatch> sessions = new ConcurrentHashMap<>();
    private final Map<String, OutputStream> inputPipes = new ConcurrentHashMap<>();
    private ThreadPoolExecutor threadPoolExecutor;
    private final KubernetesClient kubernetesClient;

    public K8sTerminalWebSocketHandler(ThreadPoolExecutor threadPoolExecutor, KubernetesClient kubernetesClient) {
        this.threadPoolExecutor = threadPoolExecutor;
        this.kubernetesClient = kubernetesClient;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 解析参数
        URI uri = session.getUri();
        MultiValueMap<String, String> params = UriComponentsBuilder.fromUri(uri).build().getQueryParams();

        String uid = params.getFirst("uid");
        String namespace = params.getFirst("namespace");
        String podName = params.getFirst("pod");
        String containerName = params.getFirst("container");

        System.out.printf("UID: %s | Namespace: %s | Pod: %s | Container: %s%n", uid, namespace, podName, containerName);

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

        ExecWatch exec = kubernetesClient.pods()
                .inNamespace(namespace)
                .withName(podName)
                .inContainer(containerName)
                .redirectingInput()
                .writingOutput(outputStream)
                .withTTY()
                .exec("sh");  // 可根据需要执行 "bash"

        sessions.put(session.getId(), exec);
        inputPipes.put(session.getId(), exec.getInput());

        threadPoolExecutor.submit(() -> {
            try {
                while (session.isOpen()) {
                    Thread.sleep(200);
                    String output = outputStream.toString(StandardCharsets.UTF_8);
                    if (!output.isEmpty()) {
                        session.sendMessage(new TextMessage(output));
                        outputStream.reset();
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        });
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        OutputStream input = inputPipes.get(session.getId());
        if (input != null) {
            input.write((message.getPayload()).getBytes(StandardCharsets.UTF_8));
            input.flush();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        ExecWatch exec = sessions.remove(session.getId());
        OutputStream input = inputPipes.remove(session.getId());

        if (exec != null) {
            exec.close();
        }
        if (input != null){
            input.close();
        }
    }
}
